package com.edata.bigdata

import com.vesoft.nebula.PropertyType
import com.vesoft.nebula.meta.ColumnTypeDef
import org.apache.spark.sql.{DataFrame, DataFrameReader, DataFrameWriter, Row}

import scala.collection.mutable
import scala.util.matching.Regex

package object nebulasource {

  implicit class NebulaReader(reader: DataFrameReader) {


    /*

    "NEBULA_CONN_GRAPH_ADDRESS" -> "192.168.36.128:9669"
    "NEBULA_CONN_GRAPH_USER" -> "root"
    "NEBULA_CONN_GRAPH_PASSWORD" -> "123"
    "NEBULA_CONN_ENABLE_SSL" -> "false"
    "NEBULA_CONN_TIMEOUT" -> "30"
    "NEBULA_CONN_RETRY" -> "2"
    "NEBULA_ACCESS_SPACE" -> "edata"
    "NEBULA_ACCESS_LABLE" -> "edata"
    "NEBULA_ACCESS_NO_COLS" -> "false"
    "NEBULA_ACCESS_RETURNS_COLS"->"COL_A:COLTYPE_A,COL_B:COLTYPE_B"
    "NEBULA_ACCESS_LIMIT"->"10"
    "NEBULA_ACCESS_NGQL"->" "
    ""
    */


    var connConfig: NebulaConnConfig = _
    var readConfig: NebulaReadConfig = _


    def nebula(connConfig: NebulaConnConfig, readConfig: NebulaReadConfig): NebulaReader = {
      this.connConfig = connConfig
      this.readConfig = readConfig
      this
    }

    def loadVertices(): DataFrame = {
      var ret = "v"
      var schema = ret
      if (!readConfig.no_col) {
        if (!readConfig.ret_col.isEmpty) {
          schema = readConfig.ret_col.mkString(",")
          ret = readConfig.ret_col.map(c => {
            if (c.equals("VID")) s"id(v)" else s"v.${readConfig.label}.${c}"
          }).mkString(",")
        }
      }
      val ngql = s"USE ${readConfig.space};MATCH (v:${readConfig.label}) RETURN ${ret}"
      val nebula_reader = reader.format(classOf[NebulaDataSourceProvider].getName)
      nebula_reader.option("graph_addr", connConfig.graph_addr)
      nebula_reader.option("user", connConfig.user)
      nebula_reader.option("password", connConfig.password)
      nebula_reader.option("graph_ssl", connConfig.graph_ssl)
      nebula_reader.option("timeout", connConfig.conn_timeout)
      nebula_reader.option("ret_cols", schema)
      nebula_reader.option("ngql", ngql)
      nebula_reader.load()
    }

    def loadEdges(): DataFrame = {
      var ret = "e"
      var schema = ret
      if (!readConfig.no_col) {
        schema = readConfig.ret_col.mkString(",")
        ret = readConfig.ret_col.map(c => {
          if (c.equals("SID")) {
            s"src(e)"
          }
          else if (c.equals("TID")) {
            s"dst(e)"
          } else {
            s"e.${c}"
          }
        }).mkString(",")
      }
      val ngql = s"USE ${readConfig.space};MATCH (v1)-[e:${readConfig.label}]-(v2) RETURN ${ret}"
      var nebula_reader = reader.format(classOf[NebulaDataSourceProvider].getName)
      nebula_reader.option("graph_addr", connConfig.graph_addr)
      nebula_reader.option("user", connConfig.user)
      nebula_reader.option("password", connConfig.password)
      nebula_reader.option("graph_ssl", connConfig.graph_ssl)
      nebula_reader.option("timeout", connConfig.conn_timeout)
      nebula_reader.option("ret_cols", schema)
      nebula_reader.option("ngql", ngql)
      nebula_reader.load()
    }

    def loadVerticesByNgql(): DataFrame = {
      val pattern = new Regex(s"([A-Za-z]+):${readConfig.label}")
      val matched = pattern.findFirstMatchIn(readConfig.ngql)
      var labelReplaceMent = ""
      matched.foreach { m =>
        labelReplaceMent = m.group(1)
      }
      var ret = labelReplaceMent
      var schema = ret
      if (!readConfig.no_col) {
        if (!readConfig.ret_col.isEmpty) {
          schema = readConfig.ret_col.mkString(",")
          ret = readConfig.ret_col.map(c => {
            if (c.equals("VID")) s"id(${labelReplaceMent})" else s"${labelReplaceMent}.${readConfig.label}.${c}"
          }).mkString(",")
        }
      }

      val ret_idx = findMatchStrIndexIgnoreCase(readConfig.ngql, "return")
      val ngql = readConfig.ngql.substring(0, ret_idx) + " RETURN " + ret
      val nebula_reader = reader.format(classOf[NebulaDataSourceProvider].getName)
      nebula_reader.option("graph_addr", connConfig.graph_addr)
      nebula_reader.option("user", connConfig.user)
      nebula_reader.option("password", connConfig.password)
      nebula_reader.option("graph_ssl", connConfig.graph_ssl)
      nebula_reader.option("timeout", connConfig.conn_timeout)
      nebula_reader.option("ret_cols", schema)
      nebula_reader.option("ngql", ngql)
      nebula_reader.load()
    }

    def loadEdgeByNgql(): DataFrame = {
      val pattern = new Regex(s"([A-Za-z]+):${readConfig.label}")
      val matched = pattern.findFirstMatchIn(readConfig.ngql)
      var labelReplaceMent = ""
      matched.foreach { m =>
        labelReplaceMent = m.group(1)
      }
      var ret = labelReplaceMent
      var schema = ret
      if (!readConfig.no_col) {
        if (!readConfig.ret_col.isEmpty) {
          schema = readConfig.ret_col.mkString(",")
          ret = readConfig.ret_col.map(c => {
            if (c.equals("SRC")) {
              s"src(${labelReplaceMent})"
            } else if (c.equals("DST")) {
              s"dst(${labelReplaceMent})"
            } else {
              s"${labelReplaceMent}.${c}"
            }
          }).mkString(",")
        }
      }
      val ret_idx = findMatchStrIndexIgnoreCase(readConfig.ngql, "return")
      val ngql = readConfig.ngql.substring(0, ret_idx) + " RETURN " + ret
      val nebula_reader = reader.format(classOf[NebulaDataSourceProvider].getName)
      nebula_reader.option("graph_addr", connConfig.graph_addr)
      nebula_reader.option("user", connConfig.user)
      nebula_reader.option("password", connConfig.password)
      nebula_reader.option("graph_ssl", connConfig.graph_ssl)
      nebula_reader.option("timeout", connConfig.conn_timeout)
      nebula_reader.option("ret_cols", schema)
      nebula_reader.option("ngql", ngql)
      nebula_reader.load()
    }



    def findMatchStrIndexIgnoreCase(str: String, substr: String): Int = {
      val upperStr = substr.toUpperCase()
      val lowerStr = substr.toLowerCase()
      val upperIdx = str.indexOf(upperStr)
      val lowerIdx = str.indexOf(lowerStr)
      if (upperIdx > lowerIdx) upperIdx else lowerIdx
    }
  }

  implicit class NebulaWriter(writer: DataFrameWriter[Row]) {
    var connConfig: NebulaConnConfig = _
    var writeConfig: NebulaWriteConfig = _

    def nebula(connConfig: NebulaConnConfig, writeConfig: NebulaWriteConfig): NebulaWriter = {
      this.connConfig = connConfig
      this.writeConfig = writeConfig
      this
    }

    def writeData(): Unit = {
      var nebula_writer = writer.format(classOf[NebulaDataSourceProvider].getName)
      nebula_writer.option("graph_addr", connConfig.graph_addr)
      nebula_writer.option("user", connConfig.user)
      nebula_writer.option("password", connConfig.password)
      nebula_writer.option("graph_ssl", connConfig.graph_ssl)
      nebula_writer.option("timeout", connConfig.conn_timeout)
      nebula_writer.option("gid", writeConfig.gid_field)
      nebula_writer.option("schema", writeConfig.schema)
      nebula_writer.option("gidAsProp", writeConfig.gidAsProp)
      //nebulaConfig.foreach { case (key, value) => nebula_writer = nebula_writer.option(key, value) }
      //PropertyType
      nebula_writer.save()
    }

    def writeVertices(): Unit = {
      val nebula_writer = writer.format(classOf[NebulaDataSourceProvider].getName)
      nebula_writer.option("graph_addr", connConfig.graph_addr)
      nebula_writer.option("user", connConfig.user)
      nebula_writer.option("password", connConfig.password)
      nebula_writer.option("graph_ssl", connConfig.graph_ssl)
      nebula_writer.option("timeout", connConfig.conn_timeout)
      nebula_writer.option("space", writeConfig.space)
      nebula_writer.option("label", writeConfig.label)
      nebula_writer.option("label_type", writeConfig.label_type)
      nebula_writer.option("gid", writeConfig.gid_field)
      nebula_writer.option("schema", writeConfig.schema)
      nebula_writer.option("gidAsProp", writeConfig.gidAsProp)
      nebula_writer.save()
    }

    def writeEdges(): Unit = {
      val nebula_writer = writer.format(classOf[NebulaDataSourceProvider].getName)
      nebula_writer.option("graph_addr", connConfig.graph_addr)
      nebula_writer.option("user", connConfig.user)
      nebula_writer.option("password", connConfig.password)
      nebula_writer.option("graph_ssl", connConfig.graph_ssl)
      nebula_writer.option("timeout", connConfig.conn_timeout)
      nebula_writer.option("space", writeConfig.space)
      nebula_writer.option("label", writeConfig.label)
      nebula_writer.option("label_type", writeConfig.label_type)
      nebula_writer.option("gid", writeConfig.gid_field)
      nebula_writer.option("schema", writeConfig.schema)
      nebula_writer.option("gidAsProp", writeConfig.gidAsProp)
      nebula_writer.save()
    }
  }


}
